package com.shuai.registry;

import cn.hutool.core.collection.ConcurrentHashSet;
import cn.hutool.json.JSONUtil;
import com.shuai.common.ServiceMetaInfo;
import com.shuai.config.RegistryConfig;
import io.etcd.jetcd.*;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.watch.WatchEvent;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

public class EtcdRegistry implements Registry {
  /** 正在监听的 key 集合 */
  private final Set<String> watchingKeySet = new ConcurrentHashSet<>();

  /** 注册中心服务缓存 */
  private final RegistryServiceCache registryServiceCache = new RegistryServiceCache();

  private Client client;

  private KV kvClient;

  /** 根节点 */
  private static final String ETCD_ROOT_PATH = "/rpc/";

  @Override
  public void init(RegistryConfig registryConfig) {
    client =
        Client.builder()
            .endpoints(registryConfig.getAddress())
            .connectTimeout(Duration.ofMillis(registryConfig.getTimeout()))
            .build();
    kvClient = client.getKVClient();
  }

  @Override
  public void register(ServiceMetaInfo serviceMetaInfo) throws Exception {
    // 创建 Lease 和 KV 客户端
    Lease leaseClient = client.getLeaseClient();

    // 创建一个 30 秒的租约
    long leaseId = leaseClient.grant(30).get().getID();

    // 设置要存储的键值对
    String registerKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
    ByteSequence key = ByteSequence.from(registerKey, StandardCharsets.UTF_8);
    ByteSequence value =
        ByteSequence.from(JSONUtil.toJsonStr(serviceMetaInfo), StandardCharsets.UTF_8);

    // 将键值对与租约关联起来，并设置过期时间
    PutOption putOption = PutOption.builder().withLeaseId(leaseId).build();
    kvClient.put(key, value, putOption).get();
  }

  @Override
  public void unRegister(ServiceMetaInfo serviceMetaInfo) {
    kvClient.delete(
        ByteSequence.from(
            ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey(), StandardCharsets.UTF_8));
  }

  public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {
    // 优先从缓存获取服务
    List<ServiceMetaInfo> cachedServiceMetaInfoList = registryServiceCache.readCache();
    if (cachedServiceMetaInfoList != null) {
      return cachedServiceMetaInfoList;
    }

    // 前缀搜索，结尾一定要加 '/'
    String searchPrefix = ETCD_ROOT_PATH + serviceKey + "/";

    try {
      // 前缀查询
      GetOption getOption = GetOption.builder().isPrefix(true).build();
      List<KeyValue> keyValues =
          kvClient
              .get(ByteSequence.from(searchPrefix, StandardCharsets.UTF_8), getOption)
              .get()
              .getKvs();
      // 解析服务信息
      List<ServiceMetaInfo> serviceMetaInfoList =
          keyValues.stream()
              .map(
                  keyValue -> {
                    String key = keyValue.getKey().toString(StandardCharsets.UTF_8);
                    // 监听 key 的变化
                    watch(key);
                    String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
                    return JSONUtil.toBean(value, ServiceMetaInfo.class);
                  })
              .collect(Collectors.toList());
      // 写入服务缓存
      registryServiceCache.writeCache(serviceMetaInfoList);
      return serviceMetaInfoList;
    } catch (Exception e) {
      throw new RuntimeException("获取服务列表失败", e);
    }
  }

  /**
   * 监听（消费端）
   *
   * @param serviceNodeKey
   */
  @Override
  public void watch(String serviceNodeKey) {
    Watch watchClient = client.getWatchClient();
    // 之前未被监听，开启监听
    boolean newWatch = watchingKeySet.add(serviceNodeKey);
    if (newWatch) {
      watchClient.watch(
          ByteSequence.from(serviceNodeKey, StandardCharsets.UTF_8),
          response -> {
            for (WatchEvent event : response.getEvents()) {
              switch (event.getEventType()) {
                // key 删除时触发
                case DELETE:
                  // 清理注册服务缓存
                  registryServiceCache.clearCache();
                  break;
                case PUT:
                default:
                  break;
              }
            }
          });
    }
  }

  @Override
  public void destroy() {
    System.out.println("当前节点下线");
    // 释放资源
    if (kvClient != null) {
      kvClient.close();
    }
    if (client != null) {
      client.close();
    }
  }
}
